package ner;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;

import org.xml.sax.Attributes;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.XMLReader;
import org.xml.sax.helpers.DefaultHandler;
import org.xml.sax.helpers.XMLReaderFactory;

class Training extends DefaultHandler {
	public static final String SEPRTR = ":";
	public static final String START_TAG = "START";
	public static final String START_WORD = "+end+";
	// mostusedNETag
	String mostUsedNETag;

	// {word, {tag,count}}
	private HashMap<String, HashMap<String, Integer>> wordTagMap = new HashMap<String, HashMap<String,Integer>>();
	private HashMap<String, HashMap<String, Integer>> tagTagMap = new HashMap<String, HashMap<String,Integer>>();
	// {tag, numberoftimesitoccurs}
	private HashMap<String, Integer> tagCount = new HashMap<String, Integer>();

	// {word, mostfrequenttag}
	HashMap<String, String> wordMap = new HashMap<String, String>();

//	private HashMap<String, HashMap<String, Integer>> tagTagMap = new HashMap<String, HashMap<String, Integer>>();
	private HashMap<String, HashMap<String, Integer>> tagGivenTagWord = new HashMap<String, HashMap<String, Integer>>();
	private HashMap<String, HashMap<String, Integer>> wordGivenTagWord = new HashMap<String, HashMap<String, Integer>>();

	private String[] skipList = {"ROOT", "s", "p", "TXT", "DOC", "DOCNO", "HL", "DD", "SO", "IN", "DATELINE", };
	private HashSet<String> skipSet = new HashSet<String>();
	static final String NONER = "NONER";
	private String currTag = START_TAG;
	private String prevTag = START_TAG;
	private String prevWord = START_WORD;
	private Double LAMBDA = 0.2;

	public Training() {
		for (String skipWord : skipList) {
			skipSet.add(skipWord);
		}
	}
	private int getValCount (HashMap<String, HashMap<String, Integer>> outerMap, String key, String value) {
		HashMap<String, Integer> innerMap = outerMap.get(key);
		if(innerMap == null) {
			return 0;
		}
		if(innerMap.containsKey(value)) {
			return innerMap.get(value);
		}
		return 0;
	}
	private void insertVal (HashMap<String, HashMap<String, Integer>> outerMap, String key, String value) {
		HashMap<String, Integer> innerMap = outerMap.get(key);
		if(innerMap == null) {
			innerMap = new HashMap<String, Integer>();
		}
		int count = 0;
		if(innerMap.containsKey(value)) {
			count = innerMap.get(value);
		}
		innerMap.put(value, ++count);
		outerMap.put(key, innerMap);
	}
	private int getTotalKeyCount (HashMap<String, HashMap<String, Integer>> outerMap, String key) {
		HashMap<String, Integer> innerMap = outerMap.get(key);
		if(innerMap == null) {
			return 0;
		}
		int total = 0;
		Collection <Integer> counts = innerMap.values();
		for (int count : counts) {
			total = total + count;
		}
		return total;
	}
	public void train(String inputFile, String outputFile) {

		try {
			/* The new file will be called raw.txt */
			FileWriter writer = new FileWriter(outputFile);
			FileReader reader = new FileReader(inputFile);

			XMLReader xmlr = XMLReaderFactory.createXMLReader();

			xmlr.setContentHandler(this);
			xmlr.setErrorHandler(this);

			BufferedReader br = new BufferedReader(reader);

			try {
				format(br, writer);
				reader.close();
			}
			catch (IOException ex) {
				throw new RuntimeException(ex);
			}

			reader = new FileReader(outputFile);

			// Parse the incoming XML file.
			xmlr.parse(new InputSource(reader));
		}
		catch (IOException ex) {
			throw new RuntimeException(ex);
		}
		catch (SAXException e) {
			throw new RuntimeException(e);
		}
	}
	public static void main(String[] args)
	{
		Training train = new Training();
		String inputFile = "data/ne_files/ne.muc6.training2.keys";
		String outFile = "data/ne_files/ne.muc6.training.output.keys";
		train.train(inputFile, outFile);
	}
	public static void format(BufferedReader br, FileWriter out)
	throws IOException
	{
		String line;

		try {
			out.write("<ROOT>" + "\n");
			while ((line = br.readLine()) != null) {
				line = line.replaceAll("&", "@amp;");
				out.write(line.trim() + "\n");
			}
			out.write("</ROOT>" + "\n");
		}
		catch (IOException ex) {
			System.err.println(ex);
		}

		out.close();
		br.close();
	}
	/*
	 * Grabs the opening SGML tag. 
	 */
	@Override
	public void startElement(String uri, String name, String qName, Attributes atts)
	{
		if(skipSet.contains(qName)) {
			currTag = NONER;
		}
		else {
			currTag = atts.getValue("TYPE");
			if(currTag == null) { 
				currTag = NONER;
			}
		}
	}
	/*
	 * Grabs the closing tag. 
	 */
	@Override
	public void endElement(String uri, String name, String qName)
	{
		currTag = NONER;
	}
	/*
	 * This prints out all the text between the tags we care about to 
	 * a file. 
	 */
	@Override
	public void characters(char ch[], int start, int length)
	{
		String text = new String(ch, start, length);
		String uText = unescapeText(text);
		String[] txts = uText.split(" |[.]|\\n");
		for (String currWord : txts) {
			if (currWord != null && !currWord.isEmpty()) {
				HashMap<String,Integer> hashMap = wordTagMap.get(currWord);
				if(hashMap == null) {
					hashMap = new HashMap<String, Integer>();
				}
				int count = 0;
				if(hashMap.containsKey(currTag)) {
					count = hashMap.get(currTag);
				}
				hashMap.put(currTag, ++count);
				wordTagMap.put(currWord, hashMap);
				
				String key = prevTag+ SEPRTR + prevWord;
				String value = currTag;
				insertVal(tagGivenTagWord, key, value);
				
				key = currTag+ SEPRTR + prevWord;
				value = currWord;
				insertVal(wordGivenTagWord, key, value);
				insertVal(tagTagMap, currTag, prevTag);
				count = 0;
				if(tagCount.containsKey(currTag)) {
					count = tagCount.get(currTag);
				}
				tagCount.put(currTag, ++count);
//				System.out.println(currWord + " -> " + currTag);
				prevTag = currTag;
				prevWord = currWord;
			}
		}
	}

	public String unescapeText(String text)
	{
		text = text.replaceAll("@amp;", "&");
		text = text.replaceAll("&MD;", "-");
		text = text.replaceAll("&AMP;", "&");
		text = text.replaceAll("&LR;", "");
		text = text.replaceAll("``|''", "\"");
		return text;
	}

	@Override
	public void endDocument() throws SAXException {
		super.endDocument();
//		System.out.println(wordTagMap);
//		System.out.println("------wordGivenTagWord-------");
//		System.out.println(wordGivenTagWord);
//		System.out.println("-----tagGivenTagWord--------");
//		System.out.println(tagGivenTagWord);
		Set<String> words = wordTagMap.keySet();
		for (String word : words) {
			HashMap<String, Integer> tagMap = wordTagMap.get(word);
			int maxCount = 0;
			String bestTag = null;
			Set<String> tagSet = tagMap.keySet();
			for (String tag : tagSet) {
				Integer count = tagMap.get(tag);
				if(maxCount < count) {
					bestTag = tag;
				}
			}

			wordMap.put(word, bestTag);
		}

		int maxCount = 0;

		Set<String> tagSet = tagCount.keySet();
//		tagSet.remove(NONER);
		for (String tag : tagSet) {
			if(tag.equalsIgnoreCase(NONER))
				continue;
			if(tagCount.get(tag)>maxCount) {
				maxCount = tagCount.get(tag);
				mostUsedNETag = tag;
			}
		}
	}

	public String getMostUsedNETag() {
		return mostUsedNETag;
	}

	public HashMap<String, HashMap<String, Integer>> getWordTagMap() {
		return wordTagMap;
	}

	public HashMap<String, String> getWordMap() {
		return wordMap;
	}
//	public HashMap<String, HashMap<String, Integer>> getTagTagMap() {
//		return tagTagMap;
//	}
	public Set<String> getTagSet() {
		return tagCount.keySet();
	}

	
	double probTGTW(String tag, String tw) {
		Double numer = (double)getValCount(tagGivenTagWord, tw, tag);
		Double denom = (double)getTotalKeyCount(tagGivenTagWord, tw);
		if (denom == 0) {
			String[] list = tw.split(SEPRTR);
			String newTag = list[0];
			double numer1 = (double)(getValCount(tagTagMap, tag, newTag)+LAMBDA) ;
			double denom1 = (double)(tagCount.get(newTag)+ LAMBDA*tagCount.size());
			return (numer1/denom1 );
		}
		double ans = (numer + LAMBDA ) / (denom + tagCount.size()*LAMBDA);
		return ans;
	}

	double probWGTW(String word, String tw) {
		Double numer = (double)getValCount(wordGivenTagWord, tw, word);
		Double denom = (double)getTotalKeyCount(wordGivenTagWord, tw);
		if (denom == 0) {
			String[] list = tw.split(SEPRTR);
			String newTag = list[0];
			double numer1 = (double)(getValCount(wordTagMap, word, newTag)+LAMBDA );
			double denom1 = (double)(tagCount.get(newTag)+ LAMBDA*wordMap.size());
			return ( numer1 / denom1 );
		}
		double ans = (numer + LAMBDA ) / (denom + wordMap.size()*LAMBDA);
		return ans;
	}
	
}